import { randomUUID } from 'node:crypto';
import type { WritableStream } from 'node:stream/web';
import type { RequestContext } from '../../di';
import { SpanType } from '../../observability';
import type { TracingContext } from '../../observability';
import type { ChunkType } from '../../stream/types';
import { ToolStream } from '../../tools/stream';
import { EMITTER_SYMBOL, STREAM_FORMAT_SYMBOL } from '../constants';
import type { DefaultExecutionEngine } from '../default';
import type { ExecuteFunction } from '../step';
import { getStepResult } from '../step';
import type {
  DefaultEngineType,
  Emitter,
  ExecutionContext,
  SerializedStepFlowEntry,
  StepFlowEntry,
  StepResult,
} from '../types';

export interface ExecuteSleepParams {
  workflowId: string;
  runId: string;
  serializedStepGraph: SerializedStepFlowEntry[];
  entry: {
    type: 'sleep';
    id: string;
    duration?: number;
    fn?: ExecuteFunction<any, any, any, any, any, DefaultEngineType>;
  };
  prevStep: StepFlowEntry;
  prevOutput: any;
  stepResults: Record<string, StepResult<any, any, any, any>>;
  resume?: {
    steps: string[];
    stepResults: Record<string, StepResult<any, any, any, any>>;
    resumePayload: any;
    resumePath: number[];
  };
  executionContext: ExecutionContext;
  emitter: Emitter;
  abortController: AbortController;
  requestContext: RequestContext;
  writableStream?: WritableStream<ChunkType>;
  tracingContext: TracingContext;
}

export async function executeSleep(engine: DefaultExecutionEngine, params: ExecuteSleepParams): Promise<void> {
  const {
    workflowId,
    runId,
    entry,
    prevOutput,
    stepResults,
    emitter,
    abortController,
    requestContext,
    executionContext,
    writableStream,
    tracingContext,
  } = params;

  let { duration, fn } = entry;

  const sleepSpan = tracingContext.currentSpan?.createChildSpan({
    type: SpanType.WORKFLOW_SLEEP,
    name: `sleep: ${duration ? `${duration}ms` : 'dynamic'}`,
    attributes: {
      durationMs: duration,
      sleepType: fn ? 'dynamic' : 'fixed',
    },
    tracingPolicy: engine.options?.tracingPolicy,
  });

  if (fn) {
    const stepCallId = randomUUID();
    duration = await engine.wrapDurableOperation(`workflow.${workflowId}.sleep.${entry.id}`, async () => {
      return fn({
        runId,
        workflowId,
        mastra: engine.mastra!,
        requestContext,
        inputData: prevOutput,
        state: executionContext.state,
        setState: async (state: any) => {
          executionContext.state = state;
        },
        retryCount: -1,
        tracingContext: {
          currentSpan: sleepSpan,
        },
        getInitData: () => stepResults?.input as any,
        getStepResult: getStepResult.bind(null, stepResults),
        // TODO: this function shouldn't have suspend probably?
        suspend: async (_suspendPayload: any): Promise<any> => {},
        bail: () => {},
        abort: () => {
          abortController?.abort();
        },
        [EMITTER_SYMBOL]: emitter,
        [STREAM_FORMAT_SYMBOL]: executionContext.format,
        engine: engine.getEngineContext(),
        abortSignal: abortController?.signal,
        writer: new ToolStream(
          {
            prefix: 'workflow-step',
            callId: stepCallId,
            name: 'sleep',
            runId,
          },
          writableStream,
        ),
      });
    });

    // Update sleep span with dynamic duration
    sleepSpan?.update({
      attributes: {
        durationMs: duration,
      },
    });
  }

  try {
    await engine.executeSleepDuration(!duration || duration < 0 ? 0 : duration, entry.id, workflowId);
    sleepSpan?.end();
  } catch (e) {
    sleepSpan?.error({ error: e as Error });
    throw e;
  }
}

export interface ExecuteSleepUntilParams {
  workflowId: string;
  runId: string;
  serializedStepGraph: SerializedStepFlowEntry[];
  entry: {
    type: 'sleepUntil';
    id: string;
    date?: Date;
    fn?: ExecuteFunction<any, any, any, any, any, DefaultEngineType>;
  };
  prevStep: StepFlowEntry;
  prevOutput: any;
  stepResults: Record<string, StepResult<any, any, any, any>>;
  resume?: {
    steps: string[];
    stepResults: Record<string, StepResult<any, any, any, any>>;
    resumePayload: any;
    resumePath: number[];
  };
  executionContext: ExecutionContext;
  emitter: Emitter;
  abortController: AbortController;
  requestContext: RequestContext;
  writableStream?: WritableStream<ChunkType>;
  tracingContext: TracingContext;
}

export async function executeSleepUntil(
  engine: DefaultExecutionEngine,
  params: ExecuteSleepUntilParams,
): Promise<void> {
  const {
    workflowId,
    runId,
    entry,
    prevOutput,
    stepResults,
    emitter,
    abortController,
    requestContext,
    executionContext,
    writableStream,
    tracingContext,
  } = params;

  let { date, fn } = entry;

  const sleepUntilSpan = tracingContext.currentSpan?.createChildSpan({
    type: SpanType.WORKFLOW_SLEEP,
    name: `sleepUntil: ${date ? date.toISOString() : 'dynamic'}`,
    attributes: {
      untilDate: date,
      durationMs: date ? Math.max(0, date.getTime() - Date.now()) : undefined,
      sleepType: fn ? 'dynamic' : 'fixed',
    },
    tracingPolicy: engine.options?.tracingPolicy,
  });

  if (fn) {
    const stepCallId = randomUUID();
    const dateResult = await engine.wrapDurableOperation(`workflow.${workflowId}.sleepUntil.${entry.id}`, async () => {
      return fn({
        runId,
        workflowId,
        mastra: engine.mastra!,
        requestContext,
        inputData: prevOutput,
        state: executionContext.state,
        setState: async (state: any) => {
          executionContext.state = state;
        },
        retryCount: -1,
        tracingContext: {
          currentSpan: sleepUntilSpan,
        },
        getInitData: () => stepResults?.input as any,
        getStepResult: getStepResult.bind(null, stepResults),
        // TODO: this function shouldn't have suspend probably?
        suspend: async (_suspendPayload: any): Promise<any> => {},
        bail: () => {},
        abort: () => {
          abortController?.abort();
        },
        [EMITTER_SYMBOL]: emitter,
        [STREAM_FORMAT_SYMBOL]: executionContext.format,
        engine: engine.getEngineContext(),
        abortSignal: abortController?.signal,
        writer: new ToolStream(
          {
            prefix: 'workflow-step',
            callId: stepCallId,
            name: 'sleepUntil',
            runId,
          },
          writableStream,
        ),
      });
    });
    // Ensure date is a Date object (may be serialized as string by durable execution engines)
    date = dateResult instanceof Date ? dateResult : new Date(dateResult);

    // Update sleep until span with dynamic duration
    const time = !date ? 0 : date.getTime() - Date.now();
    sleepUntilSpan?.update({
      attributes: {
        durationMs: Math.max(0, time),
      },
    });
  }

  if (!date) {
    sleepUntilSpan?.end();
    return;
  }

  try {
    await engine.executeSleepUntilDate(date, entry.id, workflowId);
    sleepUntilSpan?.end();
  } catch (e) {
    sleepUntilSpan?.error({ error: e as Error });
    throw e;
  }
}
